from typing import Dict

from llama_index.llms.openai import OpenAI
from llama_index.llms.openai.utils import ALL_AVAILABLE_MODELS, CHAT_MODELS

from env import Env  # 将API_KEY配置在env.py文件里面

# 定义 DeepSeek 模型及其上下文长度
DEEPSEEK_MODELS: Dict[str, int] = {
    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B": 128000,
}

ALL_AVAILABLE_MODELS.update(DEEPSEEK_MODELS)
CHAT_MODELS.update(DEEPSEEK_MODELS)
API_KEY = Env.API_KEY  # 替换为实际的API密钥


def get_siliconflow_llm(api_key=API_KEY, model='deepseek-ai/DeepSeek-R1-Distill-Llama-8B',
                        api_base='https://api.siliconflow.cn/v1', temperature=0.6, **kwargs) -> OpenAI:
    llm = OpenAI(api_key=api_key, model=model, api_base=api_base, temperature=temperature, **kwargs)
    """
        创建并返回一个配置好的OpenAI实例，用于与DeepSeek模型交互。

        Args:
            model (str): 模型名称，可切换为https://cloud.siliconflow.cn/models列表的其他模型
            temperature (float): 生成文本的随机性，默认为 0.6。
            **kwargs: 其他传递给OpenAI构造函数的参数。

        Returns:
            OpenAI: 配置好的OpenAI实例。
        """
    return llm


if __name__ == '__main__':
    try:
        response = get_siliconflow_llm().complete(
            "你好，我想要开发一个接口自动化测试平台，需要具备哪些技能，请帮我制定一个详细的学习计划")
        print(response)
    except Exception as e:
        print(f"An error occurred: {e}")
